import arrow
import re

from flask import jsonify, request, g
from app.libs.redprint import RedPrint
from app.db.models.vulndb import KbVuln
from app.cache import cache
from app.libs.score import Score
from app.libs.utility import make_cache_key_global
from app.handler.v2_users import check_user_free_pkg
from app.errors import ParameterError


api = RedPrint('vulndb')


@api.route('/kb_vulns', methods=['GET'])
@cache.cached(timeout=600)
def kb_vuln():
    pipline = [
        {
            '$match': {
                'is_plugin': True,
                'is_deleted': False,
                'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}
            }
        }, {
            '$group': {
                '_id': '$source_id',
                'vuln_type': {
                    '$first': '$vuln_type'
                },
                'name': {
                    '$first': '$name'
                },
                'source_id': {
                    '$first': '$source_id'
                },
                'severity':{
                    '$first': '$severity'
                }
            }
        }
    ]
    vuln_list = list(KbVuln.objects.aggregate(pipline))
    data = {
        "data": vuln_list,
        "message": "ok",
        "code": 200
    }
    return jsonify(data)


@api.route('/kb_vulns/vuln_types', methods=['GET'])
@cache.cached(timeout=600)
def kb_vuln_type():
    return jsonify(KbVuln.objects(source__in=['nscan', 'kscan', 'nuclei'], is_plugin=True, is_deleted=False).distinct('vuln_type'))


@api.route('/get', methods=['GET'])
@check_user_free_pkg
@cache.cached(timeout=600, make_cache_key=make_cache_key_global)
def get_vuln_info():
    request_args = request.args
    vuln_name = request_args.get('vuln_name')
    vuln_type = request_args.get('vuln_type')
    vuln_period = request_args.get('vuln_period')
    page = int(request_args.get('page', 1))
    limit = int(request_args.get('limit', 10))

    query = {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False}
    if vuln_name:
        query['name.zh_cn'] = {'$regex': re.escape(vuln_name), "$options": "$i"}
    if vuln_type:
        query['vuln_type.zh_cn'] = {'$regex': vuln_type}
    if vuln_period:
        period = vuln_period.split("_")
        query['updated_at'] = {'$gt': arrow.get(period[0]).datetime, '$lt': arrow.get(period[1]).datetime}

    data = KbVuln.objects.find(query).order_by("-updated_at").paginate(page, limit)
    total = data.total

    severity_map = Score.severityMap

    new_data = []
    for item in data.items:
        new_data.append({
            'vuln_name': item.name.zh_cn,
            'vuln_type': item.vuln_type[0].zh_cn if item.vuln_type else '',
            'level': severity_map.get(item.severity, {}).get("zh_CN", "低危"),
            'published_at': arrow.get(item.published_at).isoformat() if item.published_at else '',
            'created_at': arrow.get(item.created_at).isoformat() if item.created_at else '',
            'updated_at': arrow.get(item.updated_at).isoformat() if item.updated_at else '',
            'cnnvd': item.cnnvd,
            'cve': item.cve,
            'owasp': item.owasp,
            'description': item.description.zh_cn,
            'impact': item.impact.zh_cn if item.impact else '',
            'recommendation': item.recommendation.zh_cn if item.recommendation else '',
            'reference': item.reference,
            'zoomeye_dork': item.zoomeye_dork,
            'zoomeye_count': item.zoomeye_count
        })

    results = {'data': new_data, 'page': page, 'count': total}
    return jsonify({'code': 200, 'msg': 'ok', 'results': results})


@api.route('/get_vuln_type', methods=['GET'])

@cache.cached(timeout=600, make_cache_key=make_cache_key_global)
def get_vuln_type():
    vuln_type = request.args.get('vuln_type')

    query = {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False}
    if vuln_type:
        query = {'vuln_type.zh_cn': {'$regex': vuln_type}, 'is_deleted': False}
    pipeline = [
        {'$match': query},
        {'$group': {'_id': "$vuln_type.zh_cn"}},
    ]
    # data = tuple(KbVuln.objects.aggregate(pipeline))
    data = KbVuln.objects.find(query).distinct("vuln_type.zh_cn")

    results = {
        'vuln_type': data
    }
    return jsonify({'code': 200, 'msg': 'ok', 'results': results})


@api.route('/kb_vulns/network_vuln_types', methods=['GET'])
@cache.cached(timeout=600)
def kb_network_vuln_type():
    data = [d.to_dict() for d in KbVuln.objects(event_name="network_vuln", is_deleted=False).distinct('vuln_type')]
    return jsonify(data)
